/**
 * GrfPack.cpp - This file is part of GRFPack
 * Copyright (C) 2012 FatalEror / Dean Bahari
 *
 * GRFPack is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * GRFPack is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License in main folder for more details.
 */

#include "GrfPack.hpp"

GrfFile::GrfFile(const std::string& fileName) : m_grfTableData(NULL), m_compressionLevel(Z_DEFAULT_COMPRESSION)
{
  memset(&m_grfHeader, 0, sizeof(m_grfHeader));
  this->load(fileName);
}

GrfFile::GrfFile() : m_grfTableData(NULL)
{
  memset(&m_grfHeader, 0, sizeof(m_grfHeader));
}

GrfFile::~GrfFile()
{
  this->close();
}

int GrfFile::create(const std::string& fileName)
{
  // Check whether file is exists or not
  m_fileHandle.open(fileName.c_str(), std::ios::in);
  if (m_fileHandle.good()) {
    m_fileHandle.close();
    return -2;
  }

  // Create new file
  m_fileHandle.open(fileName.c_str(), std::ios::out | std::ios::binary);
  if (!m_fileHandle.is_open())
    return -1;

  memcpy(m_grfHeader.signature, MAGIC_KEY, sizeof(m_grfHeader.signature));
  memset(m_grfHeader.allowEncryption, 0, sizeof(m_grfHeader.allowEncryption));
  m_grfHeader.fileTableOffset = 0;
  m_grfHeader.number1 = 0;
  m_grfHeader.number2 = 7;
  m_grfHeader.version = 0x200;
  m_fileHandle.write(reinterpret_cast<char*>(&m_grfHeader), sizeof(m_grfHeader));

  uint8_t tableData[2] = { 0x00, 0x00 };
  m_fileHandle.write(reinterpret_cast<char*>(tableData), sizeof(tableData));

  m_fileHandle.close();
  if (this->load(fileName) == 0)
    return 0;

  return -1;
}

void GrfFile::save()
{
  if (!m_fileHandle.is_open() || !m_tempHandle.is_open())
    return;

  // Merge data from temp file
  m_tempHandle.seekg(0, std::ios::end);
  uint32_t tempSize = static_cast<uint32_t>(m_tempHandle.tellg());
  uint8_t *tempData = new uint8_t[tempSize];
  m_tempHandle.seekg(0, std::ios::beg);
  m_tempHandle.read(reinterpret_cast<char*>(tempData), tempSize);
  m_fileHandle.seekp(sizeof(m_grfHeader) + m_grfHeader.fileTableOffset, std::ios::beg);
  m_fileHandle.write(reinterpret_cast<char*>(tempData), tempSize);
  m_grfHeader.fileTableOffset = static_cast<uint32_t>(m_fileHandle.tellp()) - sizeof(m_grfHeader);

  delete[] tempData;

  // Compress table data
  uint32_t compressedSize = compressBound(m_grfTableSize);
  uint8_t *compressedData = new uint8_t[compressedSize];
  compress2(compressedData, reinterpret_cast<uLongf*>(&compressedSize), m_grfTableData, m_grfTableSize, m_compressionLevel);
  m_fileHandle.write(reinterpret_cast<char*>(&compressedSize), sizeof(compressedSize));
  m_fileHandle.write(reinterpret_cast<char*>(&m_grfTableSize), sizeof(m_grfTableSize));
  m_fileHandle.write(reinterpret_cast<char*>(compressedData), compressedSize);

  delete[] compressedData;

  // Rewrite GRF header
  m_fileHandle.seekp(0, std::ios::beg);
  m_fileHandle.write(reinterpret_cast<char*>(&m_grfHeader), sizeof(m_grfHeader));

  // Refresh temp file
  m_tempHandle.close();
  m_tempHandle.open(m_tempName.c_str(), std::ios::in | std::ios::out | std::ios::binary | std::ios::trunc);
}

int GrfFile::load(const std::string& fileName)
{
  m_fileHandle.open(fileName.c_str(), std::ios::in | std::ios::out | std::ios::binary);
  if (!m_fileHandle.is_open())
    return -1;

  // Read GRF Header(46 bytes) and store it into m_grfHeader
  m_fileHandle.read(reinterpret_cast<char*>(&m_grfHeader), sizeof(m_grfHeader));
  // Check if GRF is valid
  if (memcmp(m_grfHeader.signature, MAGIC_KEY, sizeof(m_grfHeader.signature)) != 0) {
    m_fileHandle.close();
    return -2;
  }

  m_fileName = fileName;
  m_tempName = fileName + "~";

  // Create a temp file for storing temporary data
  m_tempHandle.open(m_tempName.c_str(), std::ios::in | std::ios::out | std::ios::binary | std::ios::trunc);

  // There's no file in GRF, no need to read table data
  if (this->getFileCount() == 0) {
    m_grfTableSize = 0;
    return 0;
  }

  // Read GRF table data
  uint32_t fileTableOffset = sizeof(m_grfHeader) + m_grfHeader.fileTableOffset;
  m_fileHandle.seekg(fileTableOffset, std::ios::beg);
  if (static_cast<uint32_t>(m_fileHandle.tellg()) != fileTableOffset) {
    m_fileHandle.close();
    return -2;
  }

  uint32_t compressedLength;

  m_fileHandle.read(reinterpret_cast<char*>(&compressedLength), sizeof(compressedLength));
  m_fileHandle.read(reinterpret_cast<char*>(&m_grfTableSize), sizeof(m_grfTableSize));

  // Uncompress table data
  uint8_t *compressedData = new uint8_t[compressedLength];
  m_grfTableData = new uint8_t[m_grfTableSize];

  m_fileHandle.read(reinterpret_cast<char*>(compressedData), compressedLength);
  uncompress(m_grfTableData, reinterpret_cast<uLongf*>(&m_grfTableSize), compressedData, compressedLength);

  this->parseTableData(m_grfTableData);

  delete[] compressedData;

  return 0;
}

void GrfFile::close()
{
  if (m_fileHandle.is_open()) {
    if (m_grfTableData)
      delete[] m_grfTableData;
    m_grfItem.clear();
    m_fileHandle.close();
    m_tempHandle.close();

    std::remove(m_tempName.c_str());
    for (uint32_t i = 0; i < m_grfItem.size(); i++)
      this->getItem(i)->free();
  }
}

int GrfFile::insert(const std::string& srcPath, const std::string& destPath)
{
  std::fstream fileHandle(srcPath.c_str(), std::ios::in | std::ios::binary);
  if (!fileHandle.good())
    return -1;

  fileHandle.seekg(0, std::ios::end);
  uint32_t fileSize = static_cast<uint32_t>(fileHandle.tellg());
  fileHandle.seekg(0, std::ios::beg);

  // File have no content
  if (fileSize == 0) {
    fileHandle.close();
    return -1;
  }

  // Store into buffer, compress it and store into temp file
  uint32_t offset = static_cast<uint32_t>(m_grfHeader.fileTableOffset + m_tempHandle.tellp());
  uint8_t *dataBuffer = new uint8_t[fileSize];
  fileHandle.read(reinterpret_cast<char*>(dataBuffer), fileSize);
  uint32_t compressedSize = compressBound(fileSize);
  uint8_t *dataCompressed = new uint8_t[compressedSize];
  compress2(dataCompressed, reinterpret_cast<uLongf*>(&compressedSize), dataBuffer, fileSize, m_compressionLevel);
  m_tempHandle.write(reinterpret_cast<char*>(dataCompressed), compressedSize);

  delete[] dataBuffer;
  delete[] dataCompressed;

  fileHandle.close();

  // Insert new item into table data
  GrfItem newItem;
  newItem.m_parent = this;
  newItem.m_offset = m_grfTableSize;
  m_grfItem.push_back(newItem);

  uint32_t len = destPath.length();
  uint32_t newTableSize = len + 18;
  uint8_t *newTableItem = new uint8_t[newTableSize];
  // File Path
  memcpy(newTableItem, destPath.c_str(), len);
  newTableItem[len] = 0x00;
  // Compressed Length
  newTableItem[len +  4] = static_cast<uint8_t>(compressedSize >> 24);
  newTableItem[len +  3] = static_cast<uint8_t>((compressedSize&0xFF0000) >> 16);
  newTableItem[len +  2] = static_cast<uint8_t>((compressedSize&0xFF00) >> 8);
  newTableItem[len +  1] = static_cast<uint8_t>(compressedSize&0xFF);
  // Compressed Length Aligned
  newTableItem[len +  8] = static_cast<uint8_t>(compressedSize >> 24);
  newTableItem[len +  7] = static_cast<uint8_t>((compressedSize&0xFF0000) >> 16);
  newTableItem[len +  6] = static_cast<uint8_t>((compressedSize&0xFF00) >> 8);
  newTableItem[len +  5] = static_cast<uint8_t>(compressedSize&0xFF);
  // Uncompressed Length
  newTableItem[len + 12] = static_cast<uint8_t>(fileSize >> 24);
  newTableItem[len + 11] = static_cast<uint8_t>((fileSize&0xFF0000) >> 16);
  newTableItem[len + 10] = static_cast<uint8_t>((fileSize&0xFF00) >> 8);
  newTableItem[len +  9] = static_cast<uint8_t>(fileSize&0xFF);
  // Flags
  newTableItem[len + 13] = GRF_FILE;
  // Offset
  newTableItem[len + 17] = static_cast<uint8_t>(offset >> 24);
  newTableItem[len + 16] = static_cast<uint8_t>((offset&0xFF0000) >> 16);
  newTableItem[len + 15] = static_cast<uint8_t>((offset&0xFF00) >> 8);
  newTableItem[len + 14] = static_cast<uint8_t>(offset&0xFF);

  uint8_t *newTableData = new uint8_t[newTableSize + m_grfTableSize];
  if (m_grfTableSize)
    memcpy(newTableData, m_grfTableData, m_grfTableSize);
  memcpy(newTableData + m_grfTableSize, newTableItem, newTableSize);
  delete[] m_grfTableData;
  m_grfTableData = newTableData;
  m_grfTableSize = newTableSize + m_grfTableSize;
  m_grfHeader.number2++;

  delete[] newTableItem;

  return 0;
}

// Get the offset for each items
void GrfFile::parseTableData(uint8_t *tableData)
{
  uint32_t ptr = 0;
  std::stringstream fileName;

  for (int i = 0; i < this->getFileCount(); i++) {
    GrfItem newItem;
    newItem.m_parent = this;
    newItem.m_offset = ptr;
    m_grfItem.push_back(newItem);

    while (tableData[ptr++] != '\0');
    ptr += 17;
  }
}

GrfItem::GrfItem() : m_fileData(NULL)
{

}

GrfItem::~GrfItem()
{
  this->free();
}

std::string GrfItem::getPath()
{
  std::stringstream stream;
  uint32_t ptr = m_offset;
  while (m_parent->m_grfTableData[ptr++] != 0)
    stream<<m_parent->m_grfTableData[ptr - 1];

  return stream.str();
}

std::string GrfItem::getName()
{
  std::string fullPath = this->getPath();
  int32_t lastSlash = fullPath.find_last_of('\\');
  if (!lastSlash)
    return fullPath;
  return fullPath.substr(lastSlash + 1);
}

inline uint32_t getUInt32FromCharArray(uint8_t *charArray, uint32_t ptr) {
  return charArray[ptr] + (charArray[ptr + 1] << 8) + (charArray[ptr + 2] << 16) + (charArray[ptr + 3] << 24);
}

uint8_t *GrfItem::getData()
{
  if (m_fileData)
    return m_fileData;

  uint32_t ptr = m_offset;
  while (m_parent->m_grfTableData[ptr++] != 0);

  uint32_t compressedLength = getUInt32FromCharArray(m_parent->m_grfTableData, ptr);
  ptr += 8;
  uint32_t realLength = getUInt32FromCharArray(m_parent->m_grfTableData, ptr);
  ptr += 5;
  uint32_t offset = getUInt32FromCharArray(m_parent->m_grfTableData, ptr);

  uint8_t *compressedData = new uint8_t[compressedLength];
  m_fileData = new uint8_t[realLength];

  m_parent->m_fileHandle.seekg(sizeof(m_parent->m_grfHeader) + offset, std::ios::beg);
  m_parent->m_fileHandle.read(reinterpret_cast<char*>(compressedData), compressedLength);

  uncompress(m_fileData, reinterpret_cast<uLongf*>(&realLength), compressedData, compressedLength);

  delete[] compressedData;

  return m_fileData;
}

uint32_t GrfItem::getCompressedSize()
{
  uint32_t ptr = m_offset;
  while (m_parent->m_grfTableData[ptr++] != 0);

  return getUInt32FromCharArray(m_parent->m_grfTableData, ptr);
}

uint32_t GrfItem::getRealSize()
{
  uint32_t ptr = m_offset;
  while (m_parent->m_grfTableData[ptr++] != 0);
  ptr += 8;

  return getUInt32FromCharArray(m_parent->m_grfTableData, ptr);
}

uint32_t GrfItem::getFlags()
{
  uint32_t ptr = m_offset;
  while (m_parent->m_grfTableData[ptr++] != 0);
  ptr += 12;

  return m_parent->m_grfTableData[ptr];
}

uint32_t GrfItem::getOffset()
{
  uint32_t ptr = m_offset;
  while (m_parent->m_grfTableData[ptr++] != 0);
  ptr += 13;

  return getUInt32FromCharArray(m_parent->m_grfTableData, ptr);
}

void GrfItem::save(const std::string& filePath)
{
  uint8_t *data = this->getData();
  std::fstream test(filePath.c_str(), std::ios::out | std::ios::binary | std::ios::trunc);
  test.write(reinterpret_cast<char*>(data), this->getRealSize());
  test.close();

  this->free();
}

void GrfItem::free()
{
  if (m_fileData) {
    delete[] m_fileData;
    m_fileData = NULL;
  }
}
